#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include "includes/myrbtree.h"

/**
 * 性质：
 * 1. 节点只有红黑两种颜色
 * 2. 根节点是黑色
 * 3. 叶节点（nil）是黑色
 * 4. 红色节点的子节点必须是黑色
 * 5. 对于每个节点，其到所有后代叶子节点的路径上，黑色节点数目相同
 */

void rbtree_insert_fixup(myrbtree *t, mynode *n);
void rbtree_remove_fixup(myrbtree *t, mynode *n);
void left_rotate(myrbtree *t, mynode *n);
void right_rotate(myrbtree *t, mynode *n);
void rbtree_transplant(myrbtree *t, mynode *u, mynode *v);
void do_travel(myrbtree *t, mynode *n);

myrbtree *rbtree_create() {
    printf("create rbtree start\n");
    myrbtree *t = malloc(sizeof(myrbtree));
    mynode *nil = node_create(0, NULL);
    nil->parent = NULL;
    nil->color = BLACK;
    nil->left = NULL;
    nil->right = NULL;
    t->size = 0;
    t->root = nil;
    t->nil = nil;
    printf("create rbtree finish\n");
    return t;
}

mynode *node_create(int key, void *val) {
    mynode *n = malloc(sizeof(mynode));
    n->key = key;
    n->value = val;
    return n;
}

void rbtree_insert(myrbtree *t, mynode *n) {
    mynode *i = t->root, *node_parent = t->nil;
    while(i != t->nil) {
        // printf("insert %d i: %d\n", n->key, i->key);
        node_parent = i;
        if(n->key < i->key) {
            i = i->left;
        } else {
            i = i->right;
        }
    }
    // printf("after find node parent\n");
    if(node_parent == t->nil) {
        t->root = n;
    } else if(n->key < node_parent->key) {
        node_parent->left = n;
    } else {
        node_parent->right = n;
    }
    n->parent = node_parent;
    n->left = t->nil;
    n->right = t->nil;
    n->color = RED;
    rbtree_insert_fixup(t, n);
}

mynode *rbtree_remove(myrbtree *t, mynode *n) {
    // n只有一个子节点
    rbtree_color removed_color = n->color;
    mynode *replacement; // 替代品
    if(n->left == t->nil) {
        replacement = n->right;
        rbtree_transplant(t, n, n->right);
    } else if(n->right == t->nil) {
        replacement = n->left;
        rbtree_transplant(t, n, n->left);
    } else {
        // n有两个子节点
        mynode *y = tree_minimum(t, n->right);
        replacement = y->right;
        removed_color = y->color;
        if(y != n->right) {
            rbtree_transplant(t, y, y->right);
            y->right = n->right;
            y->right->parent = y;
        }
        rbtree_transplant(t, n, y);
        y->left = n->left;
        y->left->parent = y;
        y->color = n->color;
    }
    if(removed_color == BLACK) {
        // 移除了一个黑色的需要重新维护性质
        rbtree_remove_fixup(t, replacement);
    }
}

void travel_tree(myrbtree *t) {
    do_travel(t, t->root);
}

void do_travel(myrbtree *t, mynode *n) {
    // 中序遍历
    if(n == t->nil) return;
    do_travel(t, n->left);
    printf("%d\n", n->key);
    do_travel(t, n->right);
}

/**
 * 循环不变性（循环开始之前）：
 * a. 节点n是红色
 * b. 如果n的父节点是根，则n的父节点是黑色
 * c. 只有性质2或4中的一个可能被破坏
 * 循环的目的是修复被破坏的性质
 */
void rbtree_insert_fixup(myrbtree *t, mynode *n) {
    // printf("start rbtree_insert_fixup\n");
    // 初始n为红色
    while(n->parent->color == RED) {
        // n和其父节点连续两个红色
        if(n->parent == n->parent->parent->left) {
            mynode *uncle = n->parent->parent->right;
            if(uncle->color == RED) {
                // 父亲和叔父节点都是红色，切换颜色和位置，保证其祖父向下的路径中黑色节点的数量不变
                // 同时解决了性质4被破坏的问题，同时将n设置为其祖父节点（红），保持了循环不变性
                uncle->color = BLACK;
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                n = n->parent->parent;
            } else {
                // 叔父节点是黑色，为了解决性质4破坏问题（n和其父为红），将n的父节点置黑，祖父置红，然后通过右旋将其父提升一层
                // 右旋之前可能要n的父节点左旋一次保证平衡
                if(n == n->parent->right) {
                    n = n->parent;
                    left_rotate(t, n);
                }
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                right_rotate(t, n->parent->parent);
            }
        } else {
            mynode *uncle = n->parent->parent->right;
            if(uncle->color == RED) {
                // 父亲和叔父节点都是红色，切换颜色和位置，保证其祖父向下的路径中黑色节点的数量不变
                // 同时解决了性质4被破坏的问题，同时将n设置为其祖父节点（红），保持了循环不变性
                uncle->color = BLACK;
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                n = n->parent->parent;
            } else {
                // 叔父节点是黑色，为了解决性质4破坏问题（n和其父为红），将n的父节点置黑，祖父置红，然后通过右旋将其父提升一层
                // 右旋之前可能要n的父节点左旋一次保证平衡
                if(n == n->parent->left) {
                    n = n->parent;
                    right_rotate(t, n);
                }
                n->parent->color = BLACK;
                n->parent->parent->color = RED;
                left_rotate(t, n->parent->parent);
            }
        }
    }
    t->root->color = BLACK;
}

void rbtree_remove_fixup(myrbtree *t, mynode *n) {
    while (n != t->root && n->color == BLACK)
    {
        // 目标：将n上额外的一层黑色转移到一个红节点或根节点上
        if(n == n->parent->left) {
            mynode *brother = n->parent->right;
            if(brother->color == RED) {
                // 情况1：兄弟节点是红色，根据性质4 brother的子节点和父节点只能是黑色
                // 处理后n的兄弟节点变成brother之前的左黑节点(情况1、2、3)
                left_rotate(t, n->parent);
                n->parent->color = RED;
                brother->color = BLACK;
            } else {
                if(brother->left->color == BLACK && brother->right->color == BLACK) {
                    // 情况2：n上移一层
                    brother->color = RED;
                    n = n->parent;
                } else {
                    if(brother->right->color == BLACK) {
                        // 情况3：将n的兄弟节点的右子节点转换为红色
                        brother->color = RED;
                        brother->left->color = BLACK;
                        right_rotate(t, brother);
                    } else {
                        n->parent->color = BLACK; // n额外的一层黑色转移到n的父节点上
                        brother->color = RED;
                        brother->right->color = BLACK;
                        left_rotate(t, n->parent);
                        n = t->root;
                    }
                }
            }

        }
    }
    
    n->color = BLACK;
}


void left_rotate(myrbtree *t, mynode *n) {
    mynode *r = n->right;
    if(n->parent == t->nil) {
        t->root = r;
    } else if(n == n->parent->left) {
        n->parent->left = r;
    } else {
        n->parent->right = r;
    }
    r->parent = n->parent;
    n->right = r->left;
    n->parent = r;
    r->left = n;
}

void right_rotate(myrbtree *t, mynode *n) {
    mynode *l = n->left;
    if(n->parent == t->nil) {
        t->root = l;
    } else if(n == n->parent->left) {
        n->parent->left = l;
    } else {
        n->parent->right = l;
    }
    l->parent = n->parent;
    n->left = l->right;
    n->parent = l;
    l->right = n;
}

void rbtree_transplant(myrbtree *t, mynode *u, mynode *v) {
    // 移除u并用v补位
    if(u == t->root) {
        t->root = v;
    } else if(u == u->parent->right) {
        u->parent->right = v;
    } else {
        u->parent->left = v;
    }
    v->parent = u->parent;
}

mynode *tree_minimum(myrbtree *t, mynode *n) {
    while (n->left != t->nil)
    {
        n = n->left;
    }
    return n;
}